package com.woniuxy.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.FileSystemResource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @Author: 马宇航
 * @Description: 知识库可以普通上传文件
 * @DateTime: 25/09/11/星期四 15:40
 * @Component: 成都蜗牛学苑
 **/
@RestController(value ="RAGFileController")
@RequestMapping("/vector")
public class RAGFileController {
    private static final Logger log = LoggerFactory.getLogger(RAGFileController.class);
    @Autowired
    RedisVectorStore redisVectorStore;
    @Value("${file.upload.path}")
    String uploadPath;
    /**
     * 做pdf的向量处理，指定文件名即可实现效果。
     * ChangeLog : 1. 创建 (25/09/11/星期四 15:47 [马宇航]);
     * @param fileName
     * @return java.lang.String
    */
    @GetMapping("/savePdf")
    @Tool(description = "如果是pdf文件，才做pdf的向量处理，指定文件名即可实现效果")
    public String savePdf(@ToolParam(description = "指定文件名一个或者多个，后缀必须是pdf") @RequestParam String ...fileName) {
        for (String s : fileName) {
            FileSystemResource fileSystemResource = new FileSystemResource(Paths.get(uploadPath + "\\" + s));
            //判空+后缀
            if (Objects.requireNonNull(fileSystemResource.getFilename()).endsWith(".pdf")) {
                PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(fileSystemResource,
                        PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(5)//每页做一个向量
                        .build());
                List<Document> documents = pagePdfDocumentReader.read();
                redisVectorStore.add(documents);
            }
        }
        log.info("文件上传成功");
        return "success";
    }

    @GetMapping("/saveTika")
    @Tool(description = "如果是非pdf非markdown的其他文件任意文件，无论是excel还是word等，都用tika的向量处理，指定文件名即可实现效果")
    public String saveTika(@ToolParam(description = "指定文件名一个或者多个，后缀可以是任意") @RequestParam String ...fileName) {
        List<Document> documents  = new ArrayList<>();
        for (String s : fileName) {
            FileSystemResource fileSystemResource = new FileSystemResource(Paths.get(uploadPath + "\\" + s));
            documents =  new TikaDocumentReader(fileSystemResource,
                        ExtractedTextFormatter.builder()
                                .withLeftAlignment(true)
                                .withNumberOfBottomTextLinesToDelete(1)
                                .withNumberOfTopTextLinesToDelete(1)
//                                .withNumberOfTopPagesToSkipBeforeDelete(1) //应该是一个动态的值，忽略前几页的 菜单，介绍
                                .build()).read();
        }
        log.info(documents.toString());
        //每10个文档存储一次
        List<Document> documents1 = new ArrayList<>();
        for (int i = 0; i < documents.size(); i++) {
             documents1.add(documents.get(i));
             //每5条存储一次
             if (i % 5 == 0) {
                 redisVectorStore.add(documents1);
                 documents1.clear();
             }
        }
        //最后不足5条的存储一次
        redisVectorStore.add(documents1);
        return "success保存："+fileName;
    }
    @GetMapping("/saveMarkdown")
    @Tool(description = "如果是markdown文件，才做markdown的向量处理，指定文件名即可实现效果")
    public String saveMarkdown(@ToolParam(description = "指定文件名一个或者多个，后缀必须是md") @RequestParam String ...fileName) {
        List<Document> documents = new ArrayList<>();
        for (String s : fileName) {
            FileSystemResource fileSystemResource = new FileSystemResource(Paths.get(uploadPath + "\\" + s));
            if (Objects.requireNonNull(fileSystemResource.getFilename()).endsWith(".md")) {
                MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(fileSystemResource,
                        MarkdownDocumentReaderConfig.builder()
                                .withIncludeCodeBlock(true)//包含代码块
                                .withHorizontalRuleCreateDocument(true)//水平规则创建文档
                                .withIncludeBlockquote(true)//包含引用
                                .build());
                documents = markdownDocumentReader.read();
                redisVectorStore.add(documents);
            }
        }
        //每10个文档存储一次
        List<Document> documents1 = new ArrayList<>();
        for (int i = 0; i < documents.size(); i++) {
            documents1.add(documents.get(i));
            //每5条存储一次
            if (i % 5 == 0) {
                redisVectorStore.add(documents1);
                documents1.clear();
            }
        }
        //最后不足5条的存储一次
        redisVectorStore.add(documents1);
        return "success保存："+fileName;
    }
}
